#include "HardSwish.h"

using namespace nvinfer1;
namespace
{
    const char *HARD_SWISH_PLUGIN_VERSION("1");
    const char *HARD_SWISH_PLUGIN_NAME("HardSwish");
}

__global__ void hardswishKernel(const float *input, float *output, int size)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size)
    {
        // 直接使用 __half 类型进行计算，避免转换为 float
        float x = input[idx];

        // 使用条件运算符代替 min 和 max
        float x_plus_3 = x + float(3.0f);
        float result = x * ((x_plus_3 < float(0.0f)) ? float(0.0f) : ((x_plus_3 > float(6.0f)) ? float(6.0f) : x_plus_3)) / float(6.0f);

        // 直接返回结果
        output[idx] = result;
    }
}

void hardSwishLauncher(const float *input, float *output, int size, cudaStream_t stream)
{
    const int blockSize = 256;
    const int gridSize = (size + blockSize - 1) / blockSize;
    hardswishKernel<<<gridSize, blockSize, 0, stream>>>(input, output, size);
}

HardSwishPlugin::HardSwishPlugin()
{
}

HardSwishPlugin::HardSwishPlugin(const void *data, size_t length)
{
}

int HardSwishPlugin::getNbOutputs() const noexcept
{
    return 1;
}

int HardSwishPlugin::initialize() noexcept
{
    return 0;
}

void HardSwishPlugin::terminate() noexcept {}

size_t HardSwishPlugin::getWorkspaceSize(const PluginTensorDesc *inputs, int nbInputs, const PluginTensorDesc *outputs, int nbOutputs) const TRTNOEXCEPT
{
    return 0;
};

// size_t HardSwishPlugin::getWorkspaceSize(int maxBatchSize) const noexcept { return 0; }

int HardSwishPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc,
	const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
	// void* output = outputs[0];

    // std::cout<<"dims.nbDims="<<inputDesc->dims.nbDims<<std::endl;

	size_t volume = 1;
	for (int i = 0; i < inputDesc->dims.nbDims; ++i) {
        // printf("%d ",inputDesc->dims.d[i]);
		volume *= inputDesc->dims.d[i];
	}
    // printf("\n");
    const float *input = static_cast<const float *>(inputs[0]);
    float *output = static_cast<float *>(outputs[0]);
    hardSwishLauncher(input, output, volume, stream);

    return 0;
}

// int HardSwishPlugin::enqueue(int batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) noexcept
// {
//     int volume = batchSize;
//     // std::cout << "batchSize="<<batchSize<<std::endl;
//     inputDims_.nbDims = 3;
//     inputDims_.d[0] = 64;
//     inputDims_.d[1] = 16;
//     inputDims_.d[2] = 16;

//     // 打印输入维度信息
//     std::cout << "enqueue Input Tensor Dimensions: ";
//     std::cout << inputDims_.nbDims << std::endl; // 有问题
//     std::cout << inputDims_.d[0] << std::endl;   // 有问题

//     const float *input = static_cast<const float *>(inputs[0]);
//     float *output = static_cast<float *>(outputs[0]);
//     hardSwishLauncher(input, output, volume, stream);

//     return 0;
// }

size_t HardSwishPlugin::getSerializationSize() const noexcept { return 0; }

void HardSwishPlugin::serialize(void *buffer) const noexcept {}

// void HardSwishPlugin::configurePlugin(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs,
//                                       const DataType *inputTypes, const DataType *outputTypes, const bool *inputIsBroadcast,
//                                       const bool *outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize)
// {
//     assert(nbInputs == 1);
//     assert(nbOutputs == 1);
//     inputDims_ = inputDims[0];
// }

bool HardSwishPlugin::supportsFormatCombination(int32_t pos, const nvinfer1::PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
	auto type = inOut[pos].type;
	auto format = inOut[pos].format;
	// 这个插件只支持普通的浮点数，以及NCHW输入格式
	if (type == nvinfer1::DataType::kFLOAT && format == nvinfer1::PluginFormat::kLINEAR) {
		return true;
	}
	else {
		return false;
	}
}

void HardSwishPlugin::configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out, int nbOutputs) TRTNOEXCEPT{
    auto type = in->desc.type;
    auto format = in->desc.format;
    assert(nbOutputs == 1);
    assert(type == nvinfer1::DataType::kFLOAT || type == nvinfer1::DataType::kHALF);
    assert(format == nvinfer1::PluginFormat::kLINEAR);
}


void HardSwishPlugin::configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) noexcept
{
    inputDims_ = inputDims[0]; // 保存输入节点shape
}

// Dims HardSwishPlugin::getOutputDimensions(int index, const Dims *inputs, int nbInputDims) noexcept
// {
//     assert(nbInputDims == 1);
//     inputDims_ = inputs[0];

//     return inputs[0];
// }

DimsExprs HardSwishPlugin::getOutputDimensions(int outputIndex, const DimsExprs *inputs, int nbInputs, IExprBuilder &exprBuilder) noexcept
{
    return *inputs;
}

// 判断数据的储存方式和类型是否符合预期条件，包括储存方式是否为线性，类型是否和第一个输入数据一致
// bool HardSwishPlugin::supportsFormat(DataType type, PluginFormat format) const
// {
//     return (type == DataType::kFLOAT && format == PluginFormat::kNCHW);
// }

const char *HardSwishPlugin::getPluginType() const noexcept
{
    return HARD_SWISH_PLUGIN_NAME;
}

const char *HardSwishPlugin::getPluginVersion() const noexcept
{
    return HARD_SWISH_PLUGIN_VERSION;
}

void HardSwishPlugin::destroy() noexcept
{
    delete this;
}

IPluginV2DynamicExt *HardSwishPlugin::clone() const noexcept
{
    IPluginV2DynamicExt *plugin = new HardSwishPlugin(); // 函数还没有重写完
    plugin->setPluginNamespace(mPluginNamespace.c_str());
    return plugin;
}

DataType HardSwishPlugin::getOutputDataType(int index, const DataType *inputTypes, int nbInputs) const noexcept
{
    assert(inputTypes && nbInputs > 0 && index == 0);
    return inputTypes[0];
}

// bool HardSwishPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted, int nbInputs) const noexcept
// {
//     return false; // HardSwish不执行广播
// }

// bool HardSwishPlugin::canBroadcastInputAcrossBatch(int inputIndex) const noexcept
// {
//     return false; // HardSwish不支持输入广播
// }

void HardSwishPlugin::setPluginNamespace(const char *pluginNamespace) noexcept
{
    mPluginNamespace = pluginNamespace ? pluginNamespace : "";
}

const char *HardSwishPlugin::getPluginNamespace() const noexcept
{
    return mPluginNamespace.c_str();
}

// 创建器的实现
// 静态成员初始化
PluginFieldCollection HardSwishPluginCreator::mFC{};
std::vector<PluginField> HardSwishPluginCreator::mPluginAttributes;

// 现在假设onnx中HardSwish算子有两个参数a和b，则
// 静态成员初始化（定义参数 a 和 b）
// std::vector<PluginField> HardSwishPluginCreator::mPluginAttributes = {
//     PluginField("a", nullptr, PluginFieldType::kFLOAT32, 1),  // 参数 a，类型 float，默认 nullptr
//     PluginField("b", nullptr, PluginFieldType::kFLOAT32, 1)   // 参数 b，类型 float，默认 nullptr
// };

// PluginFieldCollection HardSwishPluginCreator::mFC = {
//     static_cast<int>(mPluginAttributes.size()),  // mPluginAttributes参数数量
//     mPluginAttributes.data()                     // mPluginAttributes参数数组指针
// };

HardSwishPluginCreator::HardSwishPluginCreator()
{
    // 没有参数，这里完成初始化
    mPluginAttributes.clear(); // 插件占位参数为空
    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

const char *HardSwishPluginCreator::getPluginName() const
{
    return HARD_SWISH_PLUGIN_NAME;
}

const char *HardSwishPluginCreator::getPluginVersion() const
{
    return HARD_SWISH_PLUGIN_VERSION;
}

// 该函数是返回，而不是构造，所以静态成员的初始化放在函数外进行，并且这样只会初始化一次，提高高频调用时效率
const PluginFieldCollection *HardSwishPluginCreator::getFieldNames()
{
    return &mFC;
}

IPluginV2 *HardSwishPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc)
{
    return new HardSwishPlugin();
}

IPluginV2 *HardSwishPluginCreator::deserializePlugin(const char *name, const void *serialData, size_t serialLength)
{
    return new HardSwishPlugin(serialData, serialLength);
}

void HardSwishPluginCreator::setPluginNamespace(const char *pluginNamespace)
{
    mPluginNamespace = pluginNamespace;
}

const char *HardSwishPluginCreator::getPluginNamespace() const
{
    return mPluginNamespace.c_str();
}

// 注册插件
REGISTER_TENSORRT_PLUGIN(HardSwishPluginCreator);
